//========================================================================
//Copyright 2007-2010 David Yu dyuproject@gmail.com
//------------------------------------------------------------------------
//Licensed under the Apache License, Version 2.0 (the "License");
//you may not use this file except in compliance with the License.
//You may obtain a copy of the License at 
//http://www.apache.org/licenses/LICENSE-2.0
//Unless required by applicable law or agreed to in writing, software
//distributed under the License is distributed on an "AS IS" BASIS,
//WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//See the License for the specific language governing permissions and
//limitations under the License.
//========================================================================

/*
 * Copyright (c) 2009 Jens Scheffler (appenginefan.com)
 * 
 * Licensed under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in
 * compliance with the License. You may obtain a copy of the
 * License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in
 * writing, software distributed under the License is
 * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See
 * the License for the specific language governing
 * permissions and limitations under the License.
 */

package com.dyuproject.kvstore.gae;

import java.util.Collections;
import java.util.ConcurrentModificationException;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import com.dyuproject.kvstore.AbstractStore;
import com.dyuproject.kvstore.Mutable;
import com.dyuproject.kvstore.Serializer;
import com.dyuproject.kvstore.StoreException;
import com.google.appengine.api.datastore.Blob;
import com.google.appengine.api.datastore.DatastoreFailureException;
import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.DatastoreServiceFactory;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityNotFoundException;
import com.google.appengine.api.datastore.FetchOptions;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.KeyFactory;
import com.google.appengine.api.datastore.PreparedQuery;
import com.google.appengine.api.datastore.Query;
import com.google.appengine.api.datastore.Transaction;
import com.google.appengine.api.datastore.Query.FilterOperator;
import com.google.appengine.api.datastore.Query.SortDirection;
import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

/**
 * Google appengine based store.
 */
public final class GAEStore<T> extends AbstractStore<T>
{

    public static final String PREFIX = "gs:";

    public static final String PROPERTY = "b";

    public static final int DEFAULT_MAX_RETRIES = Integer.getInteger("gaestore.max_retries", 10);

    static byte[] getBlob(Entity entity)
    {
        return ((Blob)entity.getProperty(PROPERTY)).getBytes();
    }

    static Key stringToKey(String kind, String key)
    {
        return KeyFactory.createKey(kind, ":" + key);
    }

    static String keyToString(Key key)
    {
        return key.getName().substring(1);
    }

    final DatastoreService service;
    final String kind;
    final int maxRetries;

    /**
     * Constructor.
     * 
     * @param serializer
     *            Converts the byte arrays to objects
     * @param serviceOrNull
     *            a DatastoreService to use. If left null, the constructor fetch
     *            its own service
     * @param partition
     *            determines what &quot;partition&quot; to store the data in.
     *            Different stores must use different partitions, or unspecified
     *            behavior will occur.
     */
    public GAEStore(Serializer<T> serializer, DatastoreService serviceOrNull, String partition)
    {
        this(serializer, serviceOrNull, partition, DEFAULT_MAX_RETRIES);
    }
    
    /**
     * Constructor.
     * 
     * @param serializer
     *            Converts the byte arrays to objects
     * @param serviceOrNull
     *            a DatastoreService to use. If left null, the constructor fetch
     *            its own service
     * @param maxRetries
     *            the max retries tried on the function.
     * @param partition
     *            determines what &quot;partition&quot; to store the data in.
     *            Different stores must use different partitions, or unspecified
     *            behavior will occur.
     */
    public GAEStore(Serializer<T> serializer, DatastoreService serviceOrNull, String partition, 
            int maxRetries)
    {
        super(serializer);
        Preconditions.checkNotNull(partition);
        this.service = serviceOrNull == null ? DatastoreServiceFactory.getDatastoreService() : 
            serviceOrNull;
        this.kind = PREFIX + partition;
        this.maxRetries = maxRetries;
    }

    /**
     * Constructor.
     * 
     * @param serializer
     *            Converts the byte arrays to objects
     * @param partition
     *            determines what &quot;partition&quot; to store the data in.
     *            Different stores must use different partitions, or unspecified
     *            behavior will occur.
     */
    public GAEStore(Serializer<T> serializer, String partition)
    {
        this(serializer,null,partition);
    }
    
    public void put(String key, T value)
    {
        put(key, value, 1);
    }

    public void put(String key, T value, int retries)
    {
        Preconditions.checkNotNull(key);
        Preconditions.checkNotNull(value);
        Exception lastException = null;
        Key dbKey = stringToKey(kind, key);
        Entity entity = new Entity(dbKey);
        entity.setProperty(PROPERTY, new Blob(serializer.ser(value)));
        for (int i = 0; i < retries; i++)
        {
            Transaction t = service.beginTransaction();
            boolean success = false;
            try
            {
                service.put(t, entity);
                t.commit();
                success = true;
            }
            catch (ConcurrentModificationException e)
            {
                success = false;
                lastException = e;
            }
            catch (DatastoreFailureException e)
            {
                success = false;
                lastException = e;
            }
            finally
            {
                if (success)
                    return;

                t.rollback();
            }
        }

        throw new StoreException("Could not store data for key " + key,lastException);
    }

    public void delete(String key)
    {
        delete(key, 1);
    }

    public void delete(String key, int retries)
    {
        Preconditions.checkNotNull(key);
        Exception lastException = null;
        Key dbKey = stringToKey(kind, key);
        for (int i = 0; i < retries; i++)
        {
            Transaction t = service.beginTransaction();
            boolean success = false;
            try
            {
                service.delete(dbKey);
                t.commit();
                success = true;
            }
            catch (ConcurrentModificationException e)
            {
                success = false;
                lastException = e;
            }
            catch (DatastoreFailureException e)
            {
                success = false;
                lastException = e;
            }
            finally
            {
                if (success)
                    return;

                t.rollback();
            }
        }

        throw new StoreException("Could not delete data for key " + key,lastException);
    }

    public void delete(List<String> keys)
    {
        Preconditions.checkNotNull(keys);
        List<Key> storeKeys = Lists.newArrayListWithCapacity(keys.size());
        for (String k : keys)
            storeKeys.add(stringToKey(kind, k));

        try
        {
            service.delete(storeKeys);
        }
        catch (ConcurrentModificationException e)
        {
            throw new StoreException("Could not delete data for keys " + keys,e);
        }
        catch (DatastoreFailureException e)
        {
            throw new StoreException("Could not delete data for keys " + keys,e);
        }
    }

    public T mutate(String key, Function<? super T, ? extends T> mutator)
    {
        return mutate(key, mutator, maxRetries);
    }

    public T mutate(String key, Function<? super T, ? extends T> mutator, int retries)
    {
        Preconditions.checkNotNull(key);
        Preconditions.checkNotNull(mutator);
        Key dbKey = stringToKey(kind, key);
        Exception lastException = null;
        for (int i = 0; i < retries; i++)
        {
            Transaction t = service.beginTransaction();
            boolean success = false, newEntity = false;
            Entity entity = null;
            T previous = null;
            try
            {
                entity = service.get(t, dbKey);
                previous = serializer.deser(getBlob(entity));
            }
            catch (EntityNotFoundException e)
            {
                newEntity = true;
            }
            try
            {
                T data = mutator.apply(previous);
                
                if (data != null)
                {
                    if(newEntity)
                        entity = new Entity(dbKey);
                    
                    entity.setProperty(PROPERTY, new Blob(serializer.ser(data)));
                    service.put(t, entity);
                }
                else if (!newEntity)
                {
                    service.delete(t, dbKey);
                }
                t.commit();
                success = true;
                return data;
            }
            catch (ConcurrentModificationException e)
            {
                success = false;
                lastException = e;
            }
            catch (DatastoreFailureException e)
            {
                success = false;
                lastException = e;
            }
            finally
            {
                if(!success)
                    t.rollback();
            }
        }
        throw new StoreException("Could not store data for key " + key,lastException);
    }

    public T mutateRaw(Mutable<String> mutableKey, Function<? super byte[], ? extends T> mutator)
    {
        return mutateRaw(mutableKey, mutator, maxRetries);
    }

    public T mutateRaw(Mutable<String> mutableKey, Function<? super byte[], ? extends T> mutator, 
            int retries)
    {
        Preconditions.checkNotNull(mutableKey);
        Preconditions.checkNotNull(mutator);
        String key = mutableKey.getValue();
        Preconditions.checkNotNull(key);
        
        Key dbKey = stringToKey(kind, key);
        Exception lastException = null;
        for (int i = 0; i < retries; i++)
        {
            Transaction t = service.beginTransaction();
            boolean success = false, newEntity = false;
            Entity entity = null;
            byte[] b = null;
            try
            {
                entity = service.get(t, dbKey);
                b = getBlob(entity);
            }
            catch (EntityNotFoundException e)
            {
                newEntity = true;
            }
            try
            {
                T data = mutator.apply(b);
                String checkKey = mutableKey.getValue();
                // abort
                if(checkKey == null)
                    return data;
                
                if(checkKey != key)
                {
                    key = checkKey;
                    dbKey = stringToKey(kind, checkKey);
                    continue;
                }
                
                if (data != null)
                {
                    if (newEntity)
                        entity = new Entity(dbKey);

                    entity.setProperty(PROPERTY, new Blob(serializer.ser(data)));
                    service.put(t, entity);
                }
                else if (!newEntity)
                {
                    service.delete(t, dbKey);
                }
                t.commit();
                success = true;
                return data;
            }
            catch (ConcurrentModificationException e)
            {
                success = false;
                lastException = e;
            }
            catch (DatastoreFailureException e)
            {
                success = false;
                lastException = e;
            }
            finally
            {
                if(!success)
                    t.rollback();
            }
        }
        throw new StoreException("Could not store data for key " + key,lastException);
    }
    
    public <V> V get(String key, Serializer<V> serializer)
    {
        Preconditions.checkNotNull(key);
        try
        {
            Entity entity = service.get(stringToKey(kind, key));
            return serializer.deser(getBlob(entity));
        }
        catch (EntityNotFoundException e)
        {
            return null;
        }
    }
    
    public <V> List<V> get(List<String> keys, Serializer<V> serializer)
    {
        Preconditions.checkNotNull(keys);
        if (keys.size() == 0)
            Collections.emptyList();
        List<Key> storeKeys = Lists.newArrayListWithCapacity(keys.size());
        for (String k : keys)
        {
            if (k != null)
                storeKeys.add(stringToKey(kind, k));
        }
        Map<Key, Entity> entities = service.get(storeKeys);
        List<V> result = Lists.newArrayListWithCapacity(entities.size());
        for (Map.Entry<Key, Entity> entry : entities.entrySet())
        {
            result.add(serializer.deser(getBlob(entry.getValue())));
        }

        return result;
    }

    public <V> List<Entry<String, V>> getEntries(List<String> keys, Serializer<V> serializer)
    {
        Preconditions.checkNotNull(keys);
        if (keys.size() == 0)
            Collections.emptyList();
        List<Key> storeKeys = Lists.newArrayListWithCapacity(keys.size());
        for (String k : keys)
        {
            if (k != null)
                storeKeys.add(stringToKey(kind, k));
        }
        Map<Key, Entity> entities = service.get(storeKeys);
        List<Entry<String, V>> result = Lists.newArrayListWithCapacity(entities.size());
        for (Map.Entry<Key, Entity> entry : entities.entrySet())
        {
            result.add(Maps.immutableEntry(keyToString(entry.getKey()), 
                    serializer.deser(getBlob(entry.getValue()))));
        }

        return result;
    }

    <V> List<Entry<String, V>> scanEntries(Key start, Key end, int max, 
            SortDirection direction, Serializer<V> serializer)
    {
        Preconditions.checkNotNull(start);
        Preconditions.checkNotNull(end);
        Preconditions.checkArgument(max > -1);
        if (max == 0)
        {
            return Collections.emptyList();
        }
        Query query = new Query(kind);
        query.addFilter("__key__", FilterOperator.GREATER_THAN_OR_EQUAL, start);
        query.addFilter("__key__", FilterOperator.LESS_THAN, end);
        query.addSort("__key__", direction);
        PreparedQuery preparedQuery = service.prepare(query);
        List<Entry<String, V>> result = Lists.newArrayList();
        for (Entity entity : preparedQuery.asIterable(FetchOptions.Builder.withLimit(max)))
        {
            result.add(Maps.immutableEntry(keyToString(entity.getKey()), 
                    serializer.deser(getBlob(entity))));
        }
        return result;
    }
    
    <V> List<V> scan(Key start, Key end, int max, SortDirection direction, 
            Serializer<V> serializer)
    {
        Preconditions.checkNotNull(start);
        Preconditions.checkNotNull(end);
        Preconditions.checkArgument(max > -1);
        if (max == 0)
        {
            return Collections.emptyList();
        }
        Query query = new Query(kind);
        query.addFilter("__key__", FilterOperator.GREATER_THAN_OR_EQUAL, start);
        query.addFilter("__key__", FilterOperator.LESS_THAN, end);
        query.addSort("__key__", direction);
        PreparedQuery preparedQuery = service.prepare(query);
        List<V> result = Lists.newArrayList();
        for (Entity entity : preparedQuery.asIterable(FetchOptions.Builder.withLimit(max)))
        {
            result.add(serializer.deser(getBlob(entity)));
        }
        return result;
    }

    List<String> keyScan(Key start, Key end, int max, SortDirection direction)
    {
        Preconditions.checkNotNull(start);
        Preconditions.checkNotNull(end);
        Preconditions.checkArgument(max > -1);
        if (max == 0)
        {
            return Collections.emptyList();
        }
        Query query = new Query(kind);
        query.addFilter("__key__", FilterOperator.GREATER_THAN_OR_EQUAL, start);
        query.addFilter("__key__", FilterOperator.LESS_THAN, end);
        query.addSort("__key__", direction);
        query.setKeysOnly();
        PreparedQuery preparedQuery = service.prepare(query);
        List<String> result = Lists.newArrayList();
        for (Entity entity : preparedQuery.asIterable(FetchOptions.Builder.withLimit(max)))
        {
            result.add(keyToString(entity.getKey()));
        }
        return result;
    }

    public <V> List<Entry<String, V>> scanEntries(String start, String end, int max, 
            Serializer<V> serializer)
    {
        return scanEntries(stringToKey(kind, start), stringToKey(kind, end), max, 
                SortDirection.ASCENDING, serializer);
    }

    public <V> List<Entry<String, V>> scanEntriesReverse(String start, String end, int max, 
            Serializer<V> serializer)
    {
        return scanEntries(stringToKey(kind, start), stringToKey(kind, end), max, 
                SortDirection.DESCENDING, serializer);
    }
    
    public <V> List<V> scan(String start, String end, int max, Serializer<V> serializer)
    {
        return scan(stringToKey(kind, start), stringToKey(kind, end), max, 
                SortDirection.ASCENDING, serializer);
    }

    public <V> List<V> scanReverse(String start, String end, int max, Serializer<V> serializer)
    {
        return scan(stringToKey(kind, start), stringToKey(kind, end), max, 
                SortDirection.DESCENDING, serializer);
    }

    public List<String> keyScan(String start, String end, int max)
    {
        return keyScan(stringToKey(kind, start), stringToKey(kind, end), max, 
                SortDirection.ASCENDING);
    }

    public List<String> keyScanReverse(String start, String end, int max)
    {
        return keyScan(stringToKey(kind, start), stringToKey(kind, end), max, 
                SortDirection.DESCENDING);
    }

}
